/**
 * @file rstream.h
 *
 * Random stream classes and basic random number generators
 *
 * @author Dahua Lin
 */

#ifdef _MSC_VER
#pragma once
#endif

#ifndef SLIPP_RSTREAM_H_
#define SLIPP_RSTREAM_H_

#include <slipp/base/basic_defs.h>
#include <slipp/base/lang_utils.h>
#include <slipp/base/math_base.h>

#include <type_traits>
#include <random>

#ifndef SLIPP_DEFAULT_STD_RNG
#define SLIPP_DEFAULT_STD_RANDENG std::mt19937_64
#endif

namespace sli
{
	/******************************************************
	 *
	 *   standard random stream
	 *
	 ******************************************************/


	typedef SLIPP_DEFAULT_STD_RANDENG default_std_rand_eng;

	namespace _detail
	{
		template<class Eng, bool Is64Bit> struct rand_engine_adaptor;

		template<class Eng>
		struct rand_engine_adaptor<Eng, false>
		{
			int32_t next_i32(Eng& eng)
			{
				return static_cast<int32_t>(eng());
			}

			uint32_t next_u32(Eng& eng)
			{
				return static_cast<uint32_t>(eng());
			}
		};


		template<class Eng>
		struct rand_engine_adaptor<Eng, true>
		{
			bool use_low;
			union
			{
				typename Eng::result_type r;
				struct
				{
					int32_t ilow;
					int32_t ihigh;
				};
				struct
				{
					uint32_t ulow;
					uint32_t uhigh;
				};
			} store;

			rand_engine_adaptor() : use_low(true) { }

			SLIPP_ENSURE_INLINE int32_t next_i32(Eng& eng)
			{
				if (use_low)
				{
					store.r = eng();
					use_low = false;
					return store.ilow;
				}
				else
				{
					use_low = true;
					return store.ihigh;
				}
			}

			SLIPP_ENSURE_INLINE uint32_t next_u32(Eng& eng)
			{
				if (use_low)
				{
					store.r = eng();
					use_low = false;
					return store.ulow;
				}
				else
				{
					use_low = true;
					return store.uhigh;
				}
			}
		};


		template<class Eng>
		struct std_engine_traits
		{
			static const bool is_64bit = false;
		};

		template<typename UIntType, size_t w, size_t n, size_t m, size_t r,
			UIntType a, size_t u, UIntType d, size_t s,
			UIntType b, size_t t,
			UIntType c, size_t l, UIntType f>
		struct std_engine_traits<std::mersenne_twister_engine<UIntType, w, n, m, r, a, u, d, s, b, t, c, l, f> >
		{
			static const bool is_64bit = (w == 64);
		};
	}


	/**
	 * The random stream based on engines that conform to
	 * C++0x's random number engine interface
	 */
	template<class Eng=default_std_rand_eng>
	class std_rstream
	{
	public:
		typedef Eng engine_type;
		typedef typename engine_type::result_type seed_type;

	public:
		explicit std_rstream(const engine_type& eng = engine_type()) : m_eng(eng)
		{
		}

		void seed(const seed_type& s)
		{
			m_eng.seed(s);
		}

		// single

		int32_t rand_i32()
		{
			return m_adaptor.next_i32(m_eng);
		}

		uint32_t rand_u32()
		{
			return m_adaptor.next_u32(m_eng);
		}

		int32_t rand_i32(int32_t u)
		{
			return math::abs(rand_i32() % u);
		}

		uint32_t rand_u32(uint32_t u)
		{
			return rand_u32() % u;
		}

		float rand_f32()
		{
			return float(rand_u32()) * (1.0f / 0xFFFFFFFFu);
		}

		double rand_f64()
		{
			return double(rand_u32()) * (1.0 / 0xFFFFFFFFu);
		}

		// multiple

		void rand_i32_m(size_t n, int32_t *dst)
		{
			int32_t *end = dst + n;
			while (dst != end) *dst++ = rand_i32();
		}

		void rand_u32_m(size_t n, uint32_t *dst)
		{
			uint32_t *end = dst + n;
			while (dst != end) *dst++ = rand_u32();
		}

		void rand_i32_m(int32_t u, size_t n, int32_t *dst)
		{
			int32_t *end = dst + n;
			while (dst != end) *dst++ = rand_i32(u);
		}

		void rand_u32_m(uint32_t u, size_t n, uint32_t *dst)
		{
			uint32_t *end = dst + n;
			while (dst != end) *dst++ = rand_u32(u);
		}

		void rand_f32_m(size_t n, float *dst)
		{
			float *end = dst + n;
			while (dst != end) *dst++ = rand_f32();
		}

		void rand_f64_m(size_t n, double *dst)
		{
			double *end = dst + n;
			while (dst != end) *dst++ = rand_f64();
		}


	private:
		engine_type m_eng;
		_detail::rand_engine_adaptor<engine_type, _detail::std_engine_traits<engine_type>::is_64bit> m_adaptor;

	}; // end class std_rstream


	// a functor for use in integer random number generation

	template<class RStream>
	class randu32_functor
	{
	public:
		typedef uint32_t result_type;

		randu32_functor(RStream& rstream) : m_rstream(rstream) { }

		uint32_t operator() ()
		{
			return m_rstream.rand_u32();
		}

		uint32_t operator() (uint32_t u)
		{
			return m_rstream.rand_u32(u);
		}

	private:
		RStream& m_rstream;
	};


	template<class RStream>
	inline randu32_functor<RStream> randu32_fun(RStream& rstream)
	{
		return randu32_functor<RStream>(rstream);
	}



	/******************************************************
	 *
	 *   Generating bernoulli random numbers
	 *
	 ******************************************************/

	template<class RStream>
	class bernoulli_rng
	{
	public:
		typedef bool result_type;

		bernoulli_rng(RStream& rstream, double p)
		: m_rstream(rstream), m_u((uint32_t)(p * 0xFFFFFFFFu))
		{
		}

		SLIPP_ENSURE_INLINE bool get()
		{
			return m_rstream.rand_u32() <= m_u;
		}

		void get_m(size_t n, bool *dst)
		{
			bool *end = dst + n;
			while(dst != end) *dst++ = get();
		}

	private:
		RStream& m_rstream;
		uint32_t m_u;
	};


	/******************************************************
	 *
	 *   Generating discrete random numbers
	 *
	 ******************************************************/

	template<class RStream, typename IntType=int32_t>
	class discrete_rng
	{
	public:
		typedef IntType result_type;

		discrete_rng(RStream& rstream, IntType n, const double *w, double tw)
		: m_rstream(rstream), m_n_m1(n-1), m_w(w), m_tw(tw)
		{
		}

		IntType get()
		{
			double x = m_rstream.rand_f64() * m_tw;
			IntType i = 0;
			double u = m_w[0];

			while (u < x && i < m_n_m1)
			{
				u += m_w[++i];
			}
			if (u < x) ++i;

			return i;
		}

		void get_m(size_t n, IntType *dst)
		{
			IntType *end = dst + n;
			while(dst != end) *dst++ = get();
		}

	private:
		RStream& m_rstream;
		IntType m_n_m1;
		const double *m_w;
		double m_tw;
	};



	/******************************************************
	 *
	 *   Generating normal random numbers
	 *
	 ******************************************************/

	template<class RStream>
	class normal_rng
	{
	public:
		typedef double result_type;

		normal_rng(RStream& rstream)
		: m_rstream(rstream), m_store(0), m_is_stored(false)
		{
		}

		double get()
		{
			if (m_is_stored)
			{
				m_is_stored = false;
				return m_store;
			}
			else
			{
				double v;
				calc_pair(v, m_store);
				m_is_stored = true;
				return v;
			}
		}

		void get_m(size_t n, double *dst)
		{
			size_t m = n >> 1;

			while (m > 0)
			{
				calc_pair(dst[0], dst[1]);
				dst += 2;
				-- m;
			}

			if (n & 1)
			{
				*dst = get();
			}
		}

	private:
		SLIPP_ENSURE_INLINE void calc_pair(double& v1, double &v2)
		{
			double x, y, s;

			do{
				x = 1.0 - m_rstream.rand_f64() * 2.0;
				y = 1.0 - m_rstream.rand_f64() * 2.0;

				s = x * x + y * y;
			} while (s > 1);

			double t = math::sqrt((-2.0) * math::log(s) / s);

			v1 = x * t;
			v2 = y * t;
		}

	private:
		RStream& m_rstream;
		double m_store;
		bool m_is_stored;

	}; // end class normal_rng



	/******************************************************
	 *
	 *   Generating exponential random numbers
	 *
	 ******************************************************/

	template<class RStream>
	class exponential_rng
	{
	public:
		typedef double result_type;

		exponential_rng(RStream& rstream)
		: m_rstream(rstream)
		{
		}

		double get()
		{
			return - math::log(m_rstream.rand_f64());
		}

		void get_m(size_t n, double *dst)
		{
			double *end = dst + n;
			while(dst != end) *dst++ = get();
		}

	private:
		RStream& m_rstream;

	}; // end class exponential_rng


}

#endif
